use crate::ast::*;
use crate::error::{CompilerError, Result};
use super::super::TypeEnvironment;
use super::super::TryContext;
use super::super::TypeInference;

impl TypeInference {



    /// Infer types in a binary expression
    pub(in crate::type_check::type_inference) fn infer_binary_expression(
        &mut self,
        binary: &BinaryExpr,
        env: &mut TypeEnvironment,
    ) -> Result<Type> {
        let left_type = self.infer_expression(&binary.left, env)?;
        let right_type = self.infer_expression(&binary.right, env)?;

        match binary.operator {
            crate::ast::expr::BinaryOp::Add
            | crate::ast::expr::BinaryOp::Sub
            | crate::ast::expr::BinaryOp::Mul
            | crate::ast::expr::BinaryOp::Div
            | crate::ast::expr::BinaryOp::Mod => {
                // Arithmetic operations
                self.add_constraint(&left_type, &right_type, binary.location)?;
                Ok(left_type)
            }
            crate::ast::expr::BinaryOp::AddAssign
            | crate::ast::expr::BinaryOp::SubAssign
            | crate::ast::expr::BinaryOp::MulAssign
            | crate::ast::expr::BinaryOp::DivAssign
            | crate::ast::expr::BinaryOp::ModAssign => {
                self.add_constraint(&left_type, &right_type, binary.location)?;
                Ok(left_type)
            }
            crate::ast::expr::BinaryOp::Equal
            | crate::ast::expr::BinaryOp::NotEqual
            | crate::ast::expr::BinaryOp::Less
            | crate::ast::expr::BinaryOp::LessEqual
            | crate::ast::expr::BinaryOp::Greater
            | crate::ast::expr::BinaryOp::GreaterEqual => {
                // Comparison operations
                self.add_constraint(&left_type, &right_type, binary.location)?;
                Ok(Type::Basic(crate::ast::types::BasicType::Bool))
            }
            crate::ast::expr::BinaryOp::LogicalAnd | crate::ast::expr::BinaryOp::LogicalOr => {
                // Logical operations
                self.add_constraint(
                    &left_type,
                    &Type::Basic(crate::ast::types::BasicType::Bool),
                    binary.location,
                )?;
                self.add_constraint(
                    &right_type,
                    &Type::Basic(crate::ast::types::BasicType::Bool),
                    binary.location,
                )?;
                Ok(Type::Basic(crate::ast::types::BasicType::Bool))
            }
            _ => Err(CompilerError::type_error(
                binary.location.line,
                binary.location.column,
                "Unsupported binary operator",
            )),
        }
    }



    /// Infer types in a unary expression
    pub(in crate::type_check::type_inference) fn infer_unary_expression(
        &mut self,
        unary: &UnaryExpr,
        env: &mut TypeEnvironment,
    ) -> Result<Type> {
        let operand_type = self.infer_expression(&unary.operand, env)?;

        match unary.operator {
            crate::ast::expr::UnaryOp::Not => {
                self.add_constraint(
                    &operand_type,
                    &Type::Basic(crate::ast::types::BasicType::Bool),
                    unary.location,
                )?;
                Ok(Type::Basic(crate::ast::types::BasicType::Bool))
            }
            crate::ast::expr::UnaryOp::Neg | crate::ast::expr::UnaryOp::Pos => {
                // Numeric operations
                Ok(operand_type)
            }
            crate::ast::expr::UnaryOp::AddrOf => {
                Ok(Type::Reference(ReferenceType {
                    referenced_type: Box::new(operand_type),
                    location: unary.location,
                }))
            }
            crate::ast::expr::UnaryOp::Deref => match operand_type {
                Type::Reference(reference_type) => Ok(*reference_type.referenced_type),
                _ => Err(CompilerError::type_error(
                    unary.location.line,
                    unary.location.column,
                    "Cannot dereference non-reference type",
                )),
            },
            _ => Err(CompilerError::type_error(
                unary.location.line,
                unary.location.column,
                "Unsupported unary operator",
            )),
        }
    }

    /// Infer types for the try (`?`) operator
    pub(in crate::type_check::type_inference) fn infer_try_expression(
        &mut self,
        try_expr: &TryExpr,
        env: &mut TypeEnvironment,
    ) -> Result<Type> {
        let try_context = self.current_try_context().cloned();
        let expected_operand_type = try_context
            .as_ref()
            .map(|ctx| match ctx {
                TryContext::Option(option_ty) => Type::Option(option_ty.clone()),
                TryContext::Result(result_ty) => Type::Result(result_ty.clone()),
            });
        let operand_type = self.infer_expression_with_expected(
            &try_expr.expression,
            env,
            expected_operand_type.as_ref(),
        )?;
        match operand_type {
            Type::Option(option_type) => match try_context.clone() {
                Some(TryContext::Option(ctx_option)) => {
                    self.add_constraint(
                        &Type::Option(option_type.clone()),
                        &Type::Option(ctx_option.clone()),
                        try_expr.location,
                    )?;
                    Ok(*option_type.some_type.clone())
                }
                Some(TryContext::Result(_)) => Err(CompilerError::type_error(
                    try_expr.location.line,
                    try_expr.location.column,
                    "Cannot use `?` on an Option when the enclosing function returns Result",
                )),
                None => Err(CompilerError::type_error(
                    try_expr.location.line,
                    try_expr.location.column,
                    "The `?` operator requires the enclosing function to return Option<T> or Result<T, E>",
                )),
            },
            Type::Result(result_type) => match try_context {
                Some(TryContext::Result(ctx_result)) => {
                    self.add_constraint(
                        result_type.err_type.as_ref(),
                        ctx_result.err_type.as_ref(),
                        try_expr.location,
                    )?;
                    Ok(*result_type.ok_type.clone())
                }
                Some(TryContext::Option(_)) => Err(CompilerError::type_error(
                    try_expr.location.line,
                    try_expr.location.column,
                    "Cannot use `?` on a Result when the enclosing function returns Option",
                )),
                None => Err(CompilerError::type_error(
                    try_expr.location.line,
                    try_expr.location.column,
                    "The `?` operator requires the enclosing function to return Option<T> or Result<T, E>",
                )),
            },
            _ => Err(CompilerError::type_error(
                try_expr.location.line,
                try_expr.location.column,
                "The `?` operator can only be applied to Option or Result values",
            )),
        }
    }

    /// Infer types for the force-unwrap (`!`) operator
    pub(in crate::type_check::type_inference) fn infer_force_unwrap_expression(
        &mut self,
        force_expr: &ForceUnwrapExpr,
        env: &mut TypeEnvironment,
    ) -> Result<Type> {
        let operand_type = self.infer_expression(&force_expr.expression, env)?;
        match operand_type {
            Type::Option(option_type) => Ok(*option_type.some_type.clone()),
            Type::Result(result_type) => Ok(*result_type.ok_type.clone()),
            _ => Err(CompilerError::type_error(
                force_expr.location.line,
                force_expr.location.column,
                "The `!` operator can only be applied to Option or Result values",
            )),
        }
    }

}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::ast::expr::{Expression, PathExpr, TryExpr};
    use crate::ast::types::{BasicType, OptionType};
    use crate::error::Location;

    #[test]
    fn try_expression_allows_none_propagation_with_context() {
        let mut inference = TypeInference::new();
        let mut env = TypeEnvironment::new();

        let option_return = Type::Option(OptionType {
            some_type: Box::new(Type::Basic(BasicType::Int)),
            location: Location::new(0, 0, 0),
        });
        inference.push_try_context_from_type(Some(&option_return));

        let try_expr = TryExpr {
            expression: Box::new(Expression::Path(PathExpr {
                segments: vec!["Option".to_string(), "None".to_string()],
                location: Location::new(0, 0, 0),
            })),
            location: Location::new(0, 0, 0),
        };

        let inferred = inference
            .infer_try_expression(&try_expr, &mut env)
            .expect("type inference should succeed for None?");
        assert_eq!(inferred, Type::Basic(BasicType::Int));
    }
}
